from decision_tree.Node import *
from copy import deepcopy

class CvDBase(object):

    def __init__(self, max_depth=None, node=None):
        self.nodes, self.layers, self.roots = [], [], []
        self.max_depth = max_depth
        self.root = node
        self.feature_sets = []
        self.label_dic = {}
        self.prune_alpha = 1
        self.whether_continuous = None

    def __str__(self):
        return "CvDTree ({})".format(self.root.height)

    __repr__ = __str__

    def feed_data(self, x, continuous_rate=0.2):

        self.feature_sets = [set(dimension) for dimension in x.T]
        data_len, data_dim = x.shape

        self.whether_continuous = np.array([len(feat) >= continuous_rate * data_len for feat in self.feature_sets])
        self.root.feats = [i for i in range(x.shape[1])]
        self.root.feed_tree(self)

    def fit(self, x, y, alpha=None, sample_weight=None, eps=1e-8, cv_rate=0.2, train_only=False):

        _dic = {c: i for i, c in enumerate(set(y))}
        y = np.array([_dic[yy] for yy in y])
        self.label_dic = {value: key for key, value in _dic.items()}
        x = np.array(x)

        self.prune_alpha = alpha if alpha is not None else x.shape[1] / 2

        if not train_only and self.root.is_cart:

            _train_num = int(len(x) * (1-cv_rate))
            _indices = np.random.permutation(np.arange(len(x)))
            _train_indices = _indices[:_train_num]
            _test_indices = _indices[_train_num:]
            if sample_weight is not None:

                _train_weights = sample_weight[_train_indices]
                _test_weights = sample_weight[_test_indices]
                _train_weights /= np.sum(_train_weights)
                _test_weights /= np.sum(_test_weights)
            else:
                _train_weights = _test_weights = None
            x_train, y_train, _train_weights = x, y, sample_weight
            x_cv = y_cv = x[_test_indices], y[_test_indices]
        else:
            x_train, y_train, _train_weights = x, y, sample_weight
            x_cv = y_cv = _test_weights = None
        self.feed_data(x_train)

        self.root.fit(x_train, y_train, _train_weights, eps)

        self.prune(x_cv, y_cv, _test_weights)

    def reduce_nodes(self):
        for i in range(len(self.nodes)-1, -1, -1):
            if self.nodes[i].prunde:
                self.nodes.pop(i)

    def _update_layers(self):

        self.layers = [[] for _ in range(self.root.height)]
        self.root.update_layers()

    def _prune(self):
        self._update_layers()
        _tmp_nodes = []

        for _node_lst in self.layers[::-1]:
            for _node in _node_lst[::-1]:
                if _node.category is None:
                    _tmp_nodes.append(_node)
        _old = np.array([node.cost() + self.prune_alpha * len(node.leafs) for node in _tmp_nodes])
        _new = np.array([node.cost(pruned=True) + self.prune_alpha for node in _tmp_nodes])

        _mask = _old >= _new
        while True:

            if self.root.height == 1:
                return
            p = np.argmax(_mask)

            if _mask[p]:
                _tmp_nodes[p].prune()

                for i, node in enumerate(_tmp_nodes):
                    if node.affected:
                        _old[i] = node.cost() + self.prune_alpha * len(node.leafs)
                        _mask[i] = _old[i] >= _new[i]
                        node.affected = False

                for i in range(len(_tmp_nodes) -1, -1, -1):
                    if _tmp_nodes[i].pruned:
                        _tmp_nodes.pop(i)
                        _old = np.delete(_old, i)
                        _new = np.delete(_new, i)
                        _mask = np.delete(_mask, i)
            else:
                break
        self.reduce_nodes()

    def _cart_prune(self):

        self.root.cut_tree()
        _tmp_nodes = [node for node in self.nodes if node.category is None]
        _thresholds = np.array([node.get_threshold() for node in _tmp_nodes])
        while True:

            root_copy = deepcopy(self.root)
            self.roots.append(root_copy)
            if self.root.height == 1:
                break
            p = np.argmin(_thresholds)
            _tmp_nodes[p].prune()
            for i, node in enumerate(_tmp_nodes):

                if node.affected:
                    _thresholds[i] = node.get_threshold()
                    node.affected = False
            for i in range(len(_tmp_nodes)-1, -1, -1):

                if _tmp_nodes[i].pruned:
                    _tmp_nodes.pop(i)
                    _thresholds = np.delete(_thresholds, i)
        self.reduce_nodes()

    @staticmethod
    def acc(y, y_pred, weights):
        if weights is not None:
            return np.sum((np.array(y) == np.array(y_pred)) * weights) / len(y)
        return np.sum(np.array(y) == np.array(y_pred)) /len(y)

    def prune(self, x_cv, y_cv, weights):
        if self.root.is_cart:

            if x_cv is not None and y_cv is not None:
                self._cart_prune()
                _arg = np.argmax([CvDBase.acc(y_cv, tree.predict(x_cv), weights) for tree in self.roots])
                _tar_root = self.roots[_arg]

                self.nodes = []
                _tar_root.feed_tree(self)
                self.root = _tar_root
        else:
            self._prune()

    def view(self):
        self.root.view()

class CvDMeta(type):
    def __new__(mcs, *args, **kwargs):
        name, bases, attr = args[:3]
        _, _node = bases

        def __init__(self, whether_continuous=None, max_depth=None, node=None, **_kwargs):
            tmp_node = node if isinstance(node, CvDNode) else _node
            CvDBase.__init__(self, whether_continuous, max_depth, tmp_node(**_kwargs))
            self._name = name

        attr["__init__"] = __init__
        return type(name, bases, attr)

class ID3Tree(CvDBase, ID3Node, metaclass=CvDMeta):
    pass

class C45Tree(CvDBase, C45Node, metaclass=CvDMeta):
    pass

class CartTree(CvDBase, CartNode, metaclass=CvDMeta):
    pass